from rknn.api import RKNN
import onnx
import os

def simplify_yolov8_model():
    """简化YOLOv8模型"""
    # 使用onnx-simplifier简化模型
    os.system('python -m onnxsim best.onnx best_simplified.onnx')
    return 'best_simplified.onnx'

def convert_yolov8_to_rknn():
    # 先简化模型
    model_path = simplify_yolov8_model()
    
    if not os.path.exists(model_path):
        print("简化模型失败，使用原始模型")
        model_path = 'best.onnx'
    
    # 创建RKNN对象
    rknn = RKNN(verbose=True)
    
    try:
        # 使用最小化配置
        print('--> Configuring model')
        ret = rknn.config(
            target_platform='rk3588',
            # 移除所有不必要的参数
            # mean_values=[[0, 0, 0]],
            # std_values=[[255, 255, 255]],
            # 尝试优化配置
            optimization_level=3,
            # 启用更多优化选项
            # output_optimize=1
        )
        if ret != 0:
            raise Exception(f'Config failed: {ret}')
        
        # 加载ONNX
        print('--> Loading ONNX model')
        ret = rknn.load_onnx(model=model_path)
        if ret != 0:
            raise Exception(f'Load ONNX failed: {ret}')
        
        # 构建模型
        print('--> Building model')
        ret = rknn.build(
            do_quantization=False,
            # 添加这些优化选项
            # batch_size=1,
            # 尝试使用不同的优化级别
            # optimization_level=2
        )
        if ret != 0:
            raise Exception(f'Build failed: {ret}')
        
        # 导出模型
        print('--> Exporting RKNN model')
        ret = rknn.export_rknn('yolov8s_simplified.rknn')
        if ret != 0:
            raise Exception(f'Export failed: {ret}')
        
        print('✓ RKNN模型转换成功!')
        return True
        
    except Exception as e:
        print(f'✗ 错误: {e}')
        return False
    finally:
        rknn.release()

if __name__ == '__main__':
    success = convert_yolov8_to_rknn()
    print("最终结果:", "成功" if success else "失败")
